package org.systemsbiology.rface.hadoop.math.em;

import org.systemsbiology.util.KahanSummation;
import org.systemsbiology.util.Summator;

import java.util.ArrayList;

/**
 * Created by IntelliJ IDEA.
 * User: anorberg
 * Date: 10/6/11
 * Time: 4:36 PM
 * To change this template use File | Settings | File Templates.
 */
public class ExpectationMaximization1D {
    public static double iterate(Iterable<? extends EstimatedMean1D> means, Iterable<? extends PointWeightViewer1D> points){
        for(PointWeightViewer1D point : points){
            point.calculateWeights(means);
        }
        Summator updateage = new KahanSummation();
        int index = 0;
        for(EstimatedMean1D mean: means){
            for(PointWeightViewer1D point : points){
                point.setIndex(index);
            }
            updateage.add(mean.revise(points));
            ++index;
        }
        return updateage.total();
    }

    public static boolean converge(Iterable<? extends EstimatedMean1D> means, Iterable<? extends PointWeightViewer1D> points, double convergenceTarget, int maxIterations){
        for(int iter = 0; iter < maxIterations; ++iter){
            double waver = iterate(means, points);
            if(waver < convergenceTarget){
                return true;
            }
        }
        return false;
    }

    public static ArrayList<PointWeightViewer1D> makePoints(double[] in){
        ArrayList<PointWeightViewer1D> ret = new ArrayList<PointWeightViewer1D>(in.length);
        for(double d: in){
            ret.add(new PointWeightViewer1D(d));
        }
        return ret;
    }
}
